import numpy as np
import itertools
from typing import List, Iterable


class DigitalFourierTransformerBase:
    def __init__(self, N: int, K: int, **kwargs) -> None:
        self.K = K
        self.N = N
        self._dft = np.zeros([N, N], dtype=complex)
        self._A_K = [np.random.random([N, N]).astype(complex)
                     for i in range(K)]
        self.beta = 1

    @property
    def get_beta(self):
        return self.beta

    @staticmethod
    def frobenius(A: np.ndarray) -> float:
        return np.sum(np.square(np.abs(A)))

    @property
    def A_K(self):
        return self._A_K

    @property
    def dft(self):
        for i in range(self.N):
            for j in range(self.N):
                self._dft[i, j] = np.exp(-2j * np.pi * i * j / self.N)
        return self._dft

    @property
    def result(self):
        """
        输出最终的连乘积拟合结果
        """
        return self._matmul(self._A_K, self.K)

    def _matmul(self, A: List[np.ndarray], k: int):
        """
        矩阵连乘积，排除了第k个元素的顺序乘
        """
        try:
            temp = np.identity(self.N)
            for i in range(k):
                temp = temp @ A[i]
            for i in range(k + 1, self.K):
                temp = A[i] @ temp
            return temp
        except Exception as exc:
            print("here!")
            print(exc)
            return np.nan

    @property
    def error(self):
        return DigitalFourierTransformerBase.frobenius(
            np.abs(self.beta * self.dft - self._matmul(self._A_K, self.K))
        )/self.N

    def print(self, **kwargs):
        print(f"beta:{self.beta}")
        for i in range(self.K):
            print(f"A_{i}:\n{self._A_K[i]}")
        print(f"result:{self.result}")
        print(f"complex_degree:{self.comlexity}")

    def train(self, lr: float = 0.01, epoch: int = 1000, is_post_train: bool = False, **kwargs):
        raise NotImplementedError

    def post_train(self, store_path: str, **kwargs):
        raise NotImplementedError

    def grad(self, k: int, **kwargs):
        raise NotImplementedError

    def grad_beta(self):
        raise NotImplementedError

    def max_two(self, A_k: np.ndarray) -> np.ndarray:
        """
        找到最大和第二大的行元素
        """

        def _max_row(A: np.ndarray):
            assert A.shape.__len__() == 2
            array_blank = np.zeros_like(A)
            array_blank[range(A.shape[1]), np.argmax(np.abs(A), axis=1)] = A[
                range(A.shape[1]), np.argmax(np.abs(A), axis=1)
            ]
            return array_blank

        max_1_mat = _max_row(A_k)
        max_2_mat = _max_row(A_k - max_1_mat)
        return max_1_mat + max_2_mat

    @property
    def comlexity(self):
        def _compute_complex(A1, A2) -> int:
            def _compare_float(a, b):
                if abs(a - b) <= 1e-4:
                    return 1
                return 0

            def judge_complex(c):
                index = [0, 1, -1]
                _sum = sum(
                    itertools.starmap(
                        lambda x, y: _compare_float(c.real, x)
                        * _compare_float(c.imag, y),
                        itertools.product(index, index),
                    )
                )
                return 1 - _sum

            complex_degree = 0
            # for i in range(len(A1)):  # A1矩阵的行
            complex_degree = sum(
                map(
                    lambda A1_i: sum(
                        map(
                            lambda j: sum(
                                itertools.starmap(
                                    lambda x, y: judge_complex(
                                        x) * judge_complex(y[j]),
                                    zip(A1_i, A2),
                                )
                            ),
                            range(len(A2)),
                        )
                    ),
                    A1,
                )
            )
            return complex_degree

        temp = self._A_K[0]
        complex_degree = 0
        for A_K_i in self._A_K[1:]:
            complex_degree += _compute_complex(temp, A_K_i)
            temp = temp @ A_K_i
        return complex_degree
